import pandas as pd
import matplotlib.pyplot as plt
from utils.data_util import preprocessing
from sklearn.preprocessing import StandardScaler
from utils.log import Logger
from xgboost import XGBRegressor
import joblib
import matplotlib.ticker as mick
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

# 1、创建日志对象，用来保存日志消息
logger_obj = Logger("../", log_name="predict")  # 实例化日志对象
logger = logger_obj.get_logger()  # 调用日志方法


# 2.初始化读取文件数据
def init():
    df = preprocessing("../data/test.csv")
    return df


# 3.进行特征工程（特征选择、特征组合）
def feature_extract(df):

    logger.info("开始进行特征工程")

    # 1.基于time列获取到月份和小时
    df["month"] = pd.to_datetime(df["time"]).dt.month
    df["hour"] = pd.to_datetime(df["time"]).dt.hour

    # 2.获取昨天同一时间段的用电负荷
    # 2.1 将数据转换为字典类型 time作为key，power_load作为value
    time_load_dict = df.set_index("time")["power_load"].to_dict()

    # print(time_load_dict)

    # 2.2 基于当前日期获取到前一天同一时刻的日期
    # 注意：current_time、yesterday_time不是字符串类型，而是DateTime日期时间对象类型
    def get_yesterday(row, interval):
        # (1)获取到当前时间
        current_time = pd.to_datetime(row["time"])
        yesterday_time = current_time - pd.to_timedelta(interval)

        # (2)通过当前时刻获得前一天同一时刻电力负荷
        yesterday_time_str = yesterday_time.strftime("%Y-%m-%d %H:%M:%S")
        return time_load_dict.get(yesterday_time_str)

    # 对DataFrame中的数 据按行进行逐行处理
    df["yesterday"] = df.apply(get_yesterday, axis=1, args=("1d",))  # 间隔1天

    # 3- 获取当前时刻前1个小时、2个小时、3个小时的用电负荷
    # 实现方式一
    df["pre_1_hour"] = df.apply(get_yesterday, axis=1, args=("1h",))
    df["pre_2_hour"] = df.apply(get_yesterday, axis=1, args=("2h",))
    df["pre_3_hour"] = df.apply(get_yesterday, axis=1, args=("3h",))

    # 实现方式二
    # hours = 3
    # shift_list = [df["power_load"].shift(i) for i in range(1, hours + 1)]
    # shift_data = pd.concat(shift_list, axis=1)
    # shift_data.columns = [f"pre_{i}_hour" for i in range(1, hours + 1)]

    logger.info("完成特征工程")

    return df


# 4.进行模型预测
def predict_data(df):

    logger.info("开始进行模型预测")

    # 3.1 获得特征、目标值
    y = df["power_load"]

    df1 = df.drop(["power_load", "time"], axis=1)
    print(df1.head())
    x = df1.iloc[:, :]

    # 3.2 特征预处理：对测试数据进行标准化
    # 需要使用模型训练中的同一套标准StandardScaler
    # 加载transform
    transform: StandardScaler = joblib.load("../model/transformer_good_finally.pkl")
    x = transform.transform(x)

    # 3.3 加载好的训练模型对未知数据进行预测
    model: XGBRegressor = joblib.load("../model/power_load_good_finally.pkl")
    y_pre = model.predict(x)

    logger.info("完成模型预测")
    return y_pre


# 5.展示预测、真实的对比结果
def show_result(y_pre, init_df):
    # 创建画布
    fig = plt.figure(figsize=(40, 20), dpi=150)
    ax = fig.add_subplot()

    # 绘制预测和真实的曲线
    ax.plot(init_df["time"], y_pre, label="预测结果")  # 预测
    ax.plot(init_df["time"], init_df["power_load"], label="真实结果")  # 真实

    # x轴刻度尺美化
    # 横坐标时间若不处理太过密集，这里调大时间展示的间隔
    ax.xaxis.set_major_locator(mick.MultipleLocator(50))
    # 时间展示时旋转45度
    plt.xticks(rotation=45)

    # 其他属性设置
    plt.title("预测结果展示")
    plt.legend()

    plt.savefig("../model/pic1.png")
    plt.show()


if __name__ == '__main__':
    # 1.初始化：数据基本处理
    init_df = init()

    # 2.特征工程
    df = feature_extract(init_df)
    # print(new_df.head(30))

    # 3.模型预测
    y_pre = predict_data(df)

    # 4.结果预测
    show_result(y_pre, init_df)
